{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Research_work",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Linar23/Research_work/blob/master/Research_work.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tiMvMDkivYyq",
        "colab_type": "code",
        "outputId": "c84b3c25-76b5-4f99-a62c-13b17ae40f9f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 319
        }
      },
      "source": [
        "!unzip /content/Graph"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Archive:  /content/Graph.zip\n",
            "   creating: content/Graph/\n",
            "   creating: content/Graph/acc_train/\n",
            "  inflating: content/Graph/acc_train/events.out.tfevents.1559350338.18ff2474c310  \n",
            "   creating: content/Graph/avg_loss_train/\n",
            "  inflating: content/Graph/avg_loss_train/events.out.tfevents.1559350338.18ff2474c310  \n",
            "   creating: content/Graph/avg_mask_loss_test/\n",
            "  inflating: content/Graph/avg_mask_loss_test/events.out.tfevents.1559350357.18ff2474c310  \n",
            "   creating: content/Graph/avg_next_loss_test/\n",
            "  inflating: content/Graph/avg_next_loss_test/events.out.tfevents.1559350357.18ff2474c310  \n",
            "   creating: content/Graph/acc_test/\n",
            "  inflating: content/Graph/acc_test/events.out.tfevents.1559350357.18ff2474c310  \n",
            "   creating: content/Graph/avg_next_loss_train/\n",
            "  inflating: content/Graph/avg_next_loss_train/events.out.tfevents.1559350338.18ff2474c310  \n",
            "   creating: content/Graph/avg_mask_loss_train/\n",
            "  inflating: content/Graph/avg_mask_loss_train/events.out.tfevents.1559350338.18ff2474c310  \n",
            "   creating: content/Graph/avg_loss_test/\n",
            "  inflating: content/Graph/avg_loss_test/events.out.tfevents.1559350357.18ff2474c310  \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9dC2VcbQrs3x",
        "colab_type": "code",
        "outputId": "075e4239-c379-426a-fa13-7a7e20c65e48",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 302
        }
      },
      "source": [
        "!zip -r /content/Graph.zip /content/Graph"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "  adding: content/Graph/ (stored 0%)\n",
            "  adding: content/Graph/avg_mask_loss_train/ (stored 0%)\n",
            "  adding: content/Graph/avg_mask_loss_train/events.out.tfevents.1560472127.012d5558c992 (deflated 47%)\n",
            "  adding: content/Graph/avg_loss_test/ (stored 0%)\n",
            "  adding: content/Graph/avg_loss_test/events.out.tfevents.1560472153.012d5558c992 (deflated 43%)\n",
            "  adding: content/Graph/acc_train/ (stored 0%)\n",
            "  adding: content/Graph/acc_train/events.out.tfevents.1560472127.012d5558c992 (deflated 41%)\n",
            "  adding: content/Graph/avg_next_loss_train/ (stored 0%)\n",
            "  adding: content/Graph/avg_next_loss_train/events.out.tfevents.1560472127.012d5558c992 (deflated 47%)\n",
            "  adding: content/Graph/avg_next_loss_test/ (stored 0%)\n",
            "  adding: content/Graph/avg_next_loss_test/events.out.tfevents.1560472153.012d5558c992 (deflated 46%)\n",
            "  adding: content/Graph/avg_mask_loss_test/ (stored 0%)\n",
            "  adding: content/Graph/avg_mask_loss_test/events.out.tfevents.1560472153.012d5558c992 (deflated 46%)\n",
            "  adding: content/Graph/acc_test/ (stored 0%)\n",
            "  adding: content/Graph/acc_test/events.out.tfevents.1560472153.012d5558c992 (deflated 40%)\n",
            "  adding: content/Graph/avg_loss_train/ (stored 0%)\n",
            "  adding: content/Graph/avg_loss_train/events.out.tfevents.1560472127.012d5558c992 (deflated 44%)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l1Xtub1dEvsC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir /content/content/Graph"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "wr7axV9KcLNT",
        "outputId": "20c475ea-5680-42e6-80b9-c76d24f6e7f4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 84
        }
      },
      "source": [
        "!pip install tensorboardcolab\n",
        "from tensorboardcolab import TensorBoardColab\n",
        "tb = TensorBoardColab(graph_path='/content/Graph')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: tensorboardcolab in /usr/local/lib/python3.6/dist-packages (0.0.22)\n",
            "Wait for 8 seconds...\n",
            "TensorBoard link:\n",
            "https://7e915215.ngrok.io\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2DY73jpzL9lv",
        "colab_type": "code",
        "outputId": "98b6867e-0834-418c-82bf-3d91df3deddb",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 121
        }
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3Aietf%3Awg%3Aoauth%3A2.0%3Aoob&scope=email%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdocs.test%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdrive%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdrive.photos.readonly%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fpeopleapi.readonly&response_type=code\n",
            "\n",
            "Enter your authorization code:\n",
            "··········\n",
            "Mounted at /content/gdrive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OtJYjM03BFRK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!rm -r -f /content/Graph.zip"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NrnrVl-d-6Ro",
        "colab_type": "code",
        "outputId": "106cb4f9-8692-44db-9596-b9db4b5b661f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1012
        }
      },
      "source": [
        "import torch\n",
        "import random\n",
        "from collections import Counter\n",
        "from math import sqrt\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.optim import Adam\n",
        "from sklearn.model_selection import train_test_split\n",
        "import math\n",
        "import numpy as np\n",
        "import pickle\n",
        "\n",
        "\n",
        "class BERTDataset:\n",
        "    def __init__(self, corpus_path, vocab, seq_len):\n",
        "        \"\"\"\n",
        "        :param corpus_path:\n",
        "        :param vocab:\n",
        "        :param seq_len:\n",
        "        \"\"\"\n",
        "        self.vocab = vocab\n",
        "        self.seq_len = seq_len\n",
        "\n",
        "        self.corpus_path = corpus_path\n",
        "\n",
        "        with open(corpus_path, \"r\") as f:\n",
        "            self.lines = [line.replace(\"\\n\", \"\").split(\"\\\\t\") for line in f]\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.lines)\n",
        "\n",
        "    def __getitem__(self, item):\n",
        "        t1, t2, is_next_label = self.random_sent(item)\n",
        "        t1_random, t1_label = self.random_word(t1)\n",
        "        t2_random, t2_label = self.random_word(t2)\n",
        "\n",
        "        t1 = [self.vocab.cls_index] + t1_random + [self.vocab.sep_index]\n",
        "        t2 = t2_random + [self.vocab.sep_index]\n",
        "\n",
        "        t1_label = [self.vocab.pad_index] + t1_label + [self.vocab.pad_index]\n",
        "        t2_label = t2_label + [self.vocab.pad_index]\n",
        "\n",
        "        segment_label = ([1 for _ in range(len(t1))] + [2 for _ in range(len(t2))])[:self.seq_len]\n",
        "        bert_input = (t1 + t2)[:self.seq_len]\n",
        "        bert_label = (t1_label + t2_label)[:self.seq_len]\n",
        "\n",
        "        padding = [self.vocab.pad_index for _ in range(self.seq_len - len(bert_input))]\n",
        "\n",
        "        bert_input += padding\n",
        "        bert_label += padding\n",
        "        segment_label += padding\n",
        "\n",
        "        output = {\"bert_input\": bert_input,\n",
        "                  \"bert_label\": bert_label,\n",
        "                  \"segment_label\": segment_label,\n",
        "                  \"is_next\": is_next_label}\n",
        "\n",
        "        return {key: torch.tensor(value) for key, value in output.items()}\n",
        "\n",
        "    def random_word(self, sentence):\n",
        "        tokens = sentence.split()\n",
        "        output_label = []\n",
        "\n",
        "        for i, token in enumerate(tokens):\n",
        "            prob = random.random()\n",
        "            if prob < 0.15:\n",
        "                prob /= 0.15\n",
        "\n",
        "                if prob < 0.8:\n",
        "                    tokens[i] = self.vocab.mask_index\n",
        "                elif prob < 0.9:\n",
        "                    tokens[i] = random.randrange(len(self.vocab))\n",
        "                else:\n",
        "                    tokens[i] = self.vocab.token_to_index.get(token, self.vocab.unk_index)\n",
        "\n",
        "                output_label.append(self.vocab.token_to_index.get(token, self.vocab.unk_index))\n",
        "            else:\n",
        "                tokens[i] = self.vocab.token_to_index.get(token, self.vocab.unk_index)\n",
        "                output_label.append(self.vocab.pad_index)\n",
        "\n",
        "        return tokens, output_label\n",
        "\n",
        "    def random_sent(self, index):\n",
        "        t1, t2 = self.get_corpus_line(index)\n",
        "\n",
        "        if random.random() > 0.5:\n",
        "            return t1, t2, 1\n",
        "        else:\n",
        "            return t1, self.get_random_line(), 0\n",
        "\n",
        "    def get_corpus_line(self, item):\n",
        "        return self.lines[item][0], self.lines[item][1]\n",
        "\n",
        "    def get_random_line(self):\n",
        "        return self.lines[random.randrange(len(self.lines))][1]\n",
        "\n",
        "\n",
        "class Vocab:\n",
        "    def __init__(self, text):\n",
        "        \"\"\"\n",
        "        :param text:\n",
        "        \"\"\"\n",
        "        self.specials = [\"<pad>\", \"<unk>\", \"<sep>\", \"<cls>\", \"<mask>\"]\n",
        "\n",
        "        self.pad_index = 0\n",
        "        self.unk_index = 1\n",
        "        self.sep_index = 2\n",
        "        self.cls_index = 3\n",
        "        self.mask_index = 4\n",
        "\n",
        "        self.index_to_token = list(self.specials)\n",
        "\n",
        "        counter = Counter()\n",
        "        \n",
        "        for line in text:\n",
        "            words = line.replace(\"\\n\", \"\").replace(\"\\\\t\", \"\").split()\n",
        "\n",
        "            for word in words:\n",
        "                counter[word] += 1\n",
        "\n",
        "        words_and_frequencies = sorted(counter.items())\n",
        "\n",
        "        for word, freq in words_and_frequencies:\n",
        "            if(freq > 200):\n",
        "                self.index_to_token.append(word)\n",
        "\n",
        "        self.token_to_index = {token: i for i, token in enumerate(self.index_to_token)}\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.index_to_token)\n",
        "\n",
        "\n",
        "class ScaledDotProductAttention(nn.Module):\n",
        "    def __init__(self, d_k):\n",
        "        \"\"\"\n",
        "        :param d_k: int scaling factor\n",
        "        \"\"\"\n",
        "        super(ScaledDotProductAttention, self).__init__()\n",
        "\n",
        "        self.scaling = 1 / (sqrt(d_k))\n",
        "\n",
        "    def forward(self, q, k, v, mask):\n",
        "        \"\"\"\n",
        "        :param q:  An float tensor with shape of [b_s, seq_len, d_model / n_head]\n",
        "        :param k: An float tensor with shape of [b_s, seq_len, d_model / n_head]\n",
        "        :param v: An float tensor with shape of [b_s, seq_len, d_model / n_head]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, d_model / n_head]\n",
        "        \"\"\"\n",
        "        attention = torch.bmm(q, k.transpose(1, 2)) * self.scaling\n",
        "        attention = attention.masked_fill(mask == 0, -1e9)\n",
        "\n",
        "        attention = F.softmax(attention, dim=2)\n",
        "\n",
        "        output = torch.bmm(attention, v)\n",
        "\n",
        "        return output\n",
        "\n",
        "\n",
        "class SingleHeadAttention(nn.Module):\n",
        "    def __init__(self, d_model, d_k, d_v):\n",
        "        \"\"\"\n",
        "        :param d_model: Int\n",
        "        :param d_k: Int = d_model / n_head\n",
        "        :param d_v: Int = d_model / n_head\n",
        "        \"\"\"\n",
        "        super(SingleHeadAttention, self).__init__()\n",
        "\n",
        "        self.q_linear = nn.Linear(d_model, d_k)\n",
        "        self.k_linear = nn.Linear(d_model, d_k)\n",
        "        self.v_linear = nn.Linear(d_model, d_v)\n",
        "\n",
        "        self.attention = ScaledDotProductAttention(d_k)\n",
        "\n",
        "    def forward(self, q, k, v, mask):\n",
        "        \"\"\"\n",
        "        :param q: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        :param k: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        :param v: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, d_model / n_heads]\n",
        "        \"\"\"\n",
        "        proj_q = self.q_linear(q)\n",
        "        proj_k = self.k_linear(k)\n",
        "        proj_v = self.v_linear(v)\n",
        "\n",
        "        output = self.attention(proj_q, proj_k, proj_v, mask)\n",
        "\n",
        "        return output\n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, n_head, d_model):\n",
        "        \"\"\"\n",
        "        :param n_head: Int number of heads\n",
        "        :param d_model: Int\n",
        "        \"\"\"\n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "\n",
        "        d_v = int(d_model / n_head)\n",
        "        d_k = int(d_model / n_head)\n",
        "\n",
        "        self.attention = nn.ModuleList([SingleHeadAttention(d_model, d_k, d_v) for _ in range(n_head)])\n",
        "\n",
        "        self.Linear = nn.Linear(n_head * d_v, d_model)\n",
        "\n",
        "    def forward(self, q, k, v, mask):\n",
        "        \"\"\"\n",
        "        :param q: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        :param k: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        :param v: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, d_model]\n",
        "        \"\"\"\n",
        "        results = []\n",
        "\n",
        "        for i, single_attention in enumerate(self.attention):\n",
        "            attention_out = single_attention(q, k, v, mask)\n",
        "            results.append(attention_out)\n",
        "\n",
        "        concat = torch.cat(results, dim=2)\n",
        "\n",
        "        linear_output = self.Linear(concat)\n",
        "\n",
        "        return linear_output\n",
        "\n",
        "\n",
        "class TokenEmbedding(nn.Embedding):\n",
        "    def __init__(self, vocab_size, emb_size):\n",
        "        super().__init__(vocab_size, emb_size)\n",
        "\n",
        "\n",
        "class SegmentEmbedding(nn.Embedding):\n",
        "    def __init__(self, emb_size):\n",
        "        super().__init__(3, emb_size)\n",
        "\n",
        "\n",
        "class PositionalEmbedding(nn.Module):\n",
        "    def __init__(self, d_model, max_len=512):\n",
        "        super().__init__()\n",
        "\n",
        "        pe = torch.zeros(max_len, d_model).float()\n",
        "        pe.require_grad = False\n",
        "\n",
        "        position = torch.arange(0, max_len).float().unsqueeze(1)\n",
        "        div_term = torch.pow(10000, torch.arange(0, d_model, 2).float() / d_model)\n",
        "\n",
        "        pe[:, 0::2] = torch.sin(position / div_term)\n",
        "        pe[:, 1::2] = torch.cos(position / div_term)\n",
        "\n",
        "        self.pe = pe.unsqueeze(0)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.pe[:, :x.size(1)].to('cuda')\n",
        "\n",
        "\n",
        "class BERTEmbedding(nn.Module):\n",
        "    def __init__(self, vocab_size, emb_size):\n",
        "        \"\"\"\n",
        "        :param vocab_size: Int size of vocabulary\n",
        "        :param emb_size: Int size of embedding\n",
        "        \"\"\"\n",
        "        super(BERTEmbedding, self).__init__()\n",
        "\n",
        "        self.v_s = vocab_size\n",
        "        self.e_s = emb_size\n",
        "\n",
        "        self.token = TokenEmbedding(self.v_s, self.e_s)\n",
        "        self.segment = SegmentEmbedding(self.e_s)\n",
        "        self.position = PositionalEmbedding(self.e_s)\n",
        "\n",
        "    def forward(self, seq, segment_label):\n",
        "        \"\"\"\n",
        "        :param seq: An long tensor with shape of [b_s, seq_len]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, emb_size]\n",
        "        \"\"\"\n",
        "\n",
        "        return self.token(seq) + self.segment(segment_label) + self.position(seq)\n",
        "\n",
        "\n",
        "class GELU(nn.Module):\n",
        "    def forward(self, x):\n",
        "        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))\n",
        "\n",
        "\n",
        "class PositionWise(nn.Module):\n",
        "    def __init__(self, size, inner_size):\n",
        "        \"\"\"\n",
        "        :param size: Int input size\n",
        "        :param inner_size: Int inner size of position wise\n",
        "        \"\"\"\n",
        "        super(PositionWise, self).__init__()\n",
        "\n",
        "        self.fc = nn.Sequential(\n",
        "            nn.Linear(size, inner_size),\n",
        "            GELU(),\n",
        "            nn.Linear(inner_size, size)\n",
        "        )\n",
        "\n",
        "        self.layer_norm = nn.LayerNorm(size)\n",
        "\n",
        "    def forward(self, input):\n",
        "        \"\"\"\n",
        "        :param input: An float tensor with shape of [b_s, seq_len, emb_size]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, emb_size]\n",
        "        \"\"\"\n",
        "        residual = input\n",
        "\n",
        "        result = self.fc(input)\n",
        "\n",
        "        return self.layer_norm(result + residual)\n",
        "\n",
        "\n",
        "class Encoder(nn.Module):\n",
        "    def __init__(self, embeddings, d_model, n_heads, n_layers, vocab_s):\n",
        "        \"\"\"\n",
        "        :param embeddings: An float embeddings tensor with shape [b_s, seq_len, d_model]\n",
        "        :param d_model: Int size of input\n",
        "        :param n_heads: Int number of heads\n",
        "        :param vocab_s: Int size of vocabulary\n",
        "        \"\"\"\n",
        "        super(Encoder, self).__init__()\n",
        "\n",
        "        self.embeddings = embeddings\n",
        "        self.vocab_s = vocab_s\n",
        "\n",
        "        self.transformer_blocks = nn.ModuleList([nn.Sequential(MultiHeadAttention(n_heads, d_model), nn.LayerNorm(d_model), PositionWise(d_model, d_model * 4)) for _ in range(n_layers)])\n",
        "\n",
        "\n",
        "    def forward(self, x, segment_label):\n",
        "        \"\"\"\n",
        "        :param input: An long tensor with shape of [b_s, seq_len]\n",
        "\n",
        "        :return: An float tensor with shape of [b_s, seq_len, vocab_size]\n",
        "        \"\"\"\n",
        "\n",
        "        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1)\n",
        "        input = self.embeddings(x, segment_label)\n",
        "\n",
        "        for multi_head_block, layer_norm, position_wise in self.transformer_blocks:\n",
        "            input = layer_norm(input + multi_head_block(q=input, k=input, v=input, mask=mask))\n",
        "            input = position_wise(input)\n",
        "\n",
        "        return input\n",
        "\n",
        "\n",
        "class Model(nn.Module):\n",
        "    def __init__(self, n_heads, n_layers, vocab_size, emb_size):\n",
        "        \"\"\"\n",
        "        :param n_heads: Int number of heads\n",
        "        :param vocab_size: Int size of vocabulary\n",
        "        :param emb_size: Int embedding size\n",
        "        \"\"\"\n",
        "        super(Model, self).__init__()\n",
        "\n",
        "        self.embed = BERTEmbedding(vocab_size, emb_size)\n",
        "\n",
        "        self.d_model = self.embed.e_s\n",
        "        self.v_s = self.embed.v_s\n",
        "\n",
        "        self.encoder = Encoder(self.embed, self.d_model, n_heads, n_layers, self.v_s)\n",
        "        self.next_sentence = NextSentencePrediction(self.d_model)\n",
        "        self.mask_lm = MaskedLanguageModel(self.d_model, self.v_s)\n",
        "\n",
        "    def forward(self, x, segment_label):\n",
        "        \"\"\"\n",
        "        :param x: An float tensor with shape of [b_s, seq_len]\n",
        "        :param segment_label: An float tensor with shape of [b_s, seq_len]\n",
        "        \"\"\"\n",
        "        prediction = self.encoder(x, segment_label)\n",
        "\n",
        "        return self.next_sentence(prediction), self.mask_lm(prediction)\n",
        "\n",
        "\n",
        "class NextSentencePrediction(nn.Module):\n",
        "    def __init__(self, d_model):\n",
        "        \"\"\"\n",
        "        :param d_model: Int\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.linear = nn.Linear(d_model, 2)\n",
        "        self.softmax = nn.LogSoftmax(dim=-1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.softmax(self.linear(x[:, 0]))\n",
        "\n",
        "\n",
        "class MaskedLanguageModel(nn.Module):\n",
        "    def __init__(self, d_model, vocab_size):\n",
        "        \"\"\"\n",
        "        :param d_model: Int\n",
        "        :param vocab_size: Int size of vocabulary\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.linear = nn.Linear(d_model, vocab_size)\n",
        "        self.softmax = nn.Softmax(dim=-1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.softmax(self.linear(x))\n",
        "\n",
        "\n",
        "class AdamW(torch.optim.Optimizer):\n",
        "    \"\"\"Implements AdamW algorithm.\n",
        "    It has been proposed in `Fixing Weight Decay Regularization in Adam`_.\n",
        "    Arguments:\n",
        "        params (iterable): iterable of parameters to optimize or dicts defining\n",
        "            parameter groups\n",
        "        lr (float, optional): learning rate (default: 1e-3)\n",
        "        betas (Tuple[float, float], optional): coefficients used for computing\n",
        "            running averages of gradient and its square (default: (0.9, 0.999))\n",
        "        eps (float, optional): term added to the denominator to improve\n",
        "            numerical stability (default: 1e-8)\n",
        "        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n",
        "    .. Fixing Weight Decay Regularization in Adam:\n",
        "    https://arxiv.org/abs/1711.05101\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,\n",
        "                 weight_decay=0):\n",
        "        defaults = dict(lr=lr, betas=betas, eps=eps,\n",
        "                        weight_decay=weight_decay)\n",
        "        super(AdamW, self).__init__(params, defaults)\n",
        "\n",
        "    def step(self, closure=None):\n",
        "        \"\"\"Performs a single optimization step.\n",
        "        Arguments:\n",
        "            closure (callable, optional): A closure that reevaluates the model\n",
        "                and returns the loss.\n",
        "        \"\"\"\n",
        "        loss = None\n",
        "        if closure is not None:\n",
        "            loss = closure()\n",
        "\n",
        "        for group in self.param_groups:\n",
        "            for p in group['params']:\n",
        "                if p.grad is None:\n",
        "                    continue\n",
        "                grad = p.grad.data\n",
        "                if grad.is_sparse:\n",
        "                    raise RuntimeError('AdamW does not support sparse gradients, please consider SparseAdam instead')\n",
        "\n",
        "                state = self.state[p]\n",
        "\n",
        "                # State initialization\n",
        "                if len(state) == 0:\n",
        "                    state['step'] = 0\n",
        "                    # Exponential moving average of gradient values\n",
        "                    state['exp_avg'] = torch.zeros_like(p.data)\n",
        "                    # Exponential moving average of squared gradient values\n",
        "                    state['exp_avg_sq'] = torch.zeros_like(p.data)\n",
        "\n",
        "                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n",
        "                beta1, beta2 = group['betas']\n",
        "\n",
        "                state['step'] += 1\n",
        "\n",
        "                # according to the paper, this penalty should come after the bias correction\n",
        "                # if group['weight_decay'] != 0:\n",
        "                #     grad = grad.add(group['weight_decay'], p.data)\n",
        "\n",
        "                # Decay the first and second moment running average coefficient\n",
        "                exp_avg.mul_(beta1).add_(1 - beta1, grad)\n",
        "                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n",
        "\n",
        "                denom = exp_avg_sq.sqrt().add_(group['eps'])\n",
        "\n",
        "                bias_correction1 = 1 - beta1 ** state['step']\n",
        "                bias_correction2 = 1 - beta2 ** state['step']\n",
        "                step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1\n",
        "\n",
        "                # w = w - wd * lr * w\n",
        "                if group['weight_decay'] != 0:\n",
        "                    p.data.add_(-group['weight_decay'] * group['lr'], p.data)\n",
        "\n",
        "                # w = w - lr * w.grad\n",
        "                p.data.addcdiv_(-step_size, exp_avg, denom)\n",
        "\n",
        "                # w = w - wd * lr * w - lr * w.grad\n",
        "                # See http://www.fast.ai/2018/07/02/adam-weight-decay/\n",
        "\n",
        "        return loss\n",
        "\n",
        "\n",
        "class ScheduledOptim():\n",
        "    '''A simple wrapper class for learning rate scheduling'''\n",
        "\n",
        "    def __init__(self, optimizer, d_model, n_warmup_steps):\n",
        "        self._optimizer = optimizer\n",
        "        self.n_warmup_steps = n_warmup_steps\n",
        "        self.n_current_steps = 0\n",
        "        self.init_lr = np.power(d_model, -0.5)\n",
        "\n",
        "    def step_and_update_lr(self):\n",
        "        \"Step with the inner optimizer\"\n",
        "        self._update_learning_rate()\n",
        "        self._optimizer.step()\n",
        "\n",
        "    def zero_grad(self):\n",
        "        \"Zero out the gradients by the inner optimizer\"\n",
        "        self._optimizer.zero_grad()\n",
        "\n",
        "    def _get_lr_scale(self):\n",
        "        return np.min([\n",
        "            np.power(self.n_current_steps, -0.5),\n",
        "            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])\n",
        "\n",
        "    def _update_learning_rate(self):\n",
        "        ''' Learning rate scheduling per step '''\n",
        "\n",
        "        self.n_current_steps += 1\n",
        "        lr = self.init_lr * self._get_lr_scale()\n",
        "\n",
        "        for param_group in self._optimizer.param_groups:\n",
        "            param_group['lr'] = lr\n",
        "            \n",
        "            \n",
        "dataset_path = \"/content/gdrive/My Drive/data_small.txt\"\n",
        "\n",
        "with open(dataset_path, \"r\") as f:\n",
        "    vocab = Vocab(f)\n",
        "    \n",
        "#with open(\"/content/vocab_small.pickle\", \"wb\") as f:\n",
        "#    pickle.dump(vocab, f)\n",
        "    \n",
        "# with open(\"/content/dataset_small.pickle\", \"wb\") as f:\n",
        "\n",
        "#with open('/content/vocab_small.pickle', 'rb') as f:\n",
        "#    vocab = pickle.load(f)\n",
        "    \n",
        "#with open('/content/dataset_small.pickle', 'rb') as f:\n",
        "#   dataset = pickle.load(f)\n",
        "        \n",
        "print(len(vocab))\n",
        "\n",
        "seq_len = 64\n",
        "emb_size = 94\n",
        "epochs = 10\n",
        "n_layers = 4\n",
        "\n",
        "dataset = BERTDataset(dataset_path, vocab, seq_len)\n",
        "\n",
        "train, test = train_test_split(dataset, test_size=0.2)\n",
        "\n",
        "train_data_loader = DataLoader(train, batch_size=128, shuffle=True)\n",
        "test_data_loader = DataLoader(test, batch_size=128)\n",
        "#model = torch.load(\"/content/model_epoch_7.pth\").to('cuda')\n",
        "model = Model(3, n_layers, len(vocab), emb_size).to('cuda')\n",
        "masked_criterion = nn.CrossEntropyLoss(ignore_index=0)\n",
        "next_criterion = nn.CrossEntropyLoss()\n",
        "optim = AdamW(model.parameters(), lr=1e-5, weight_decay=0.01)\n",
        "optim_schedule = ScheduledOptim(optim, emb_size, n_warmup_steps=10000)\n",
        "\n",
        "train_loss = []\n",
        "test_loss = []\n",
        "train_acc = []\n",
        "test_acc = []\n",
        "\n",
        "#writer = SummaryWriter()\n",
        "\n",
        "for l in range(epochs):\n",
        "    print(\"Train\")\n",
        "\n",
        "    avg_loss = 0.0\n",
        "    avg_next_loss = 0.0\n",
        "    avg_mask_loss = 0.0\n",
        "    total_correct = 0\n",
        "    total_element = 0\n",
        "\n",
        "    for i, data in enumerate(train_data_loader):\n",
        "        #print(i / len(train_data_loader))\n",
        "        data = {key: value.to('cuda') for key, value in data.items()}\n",
        "\n",
        "        next_sent_output, mask_lm_output = model.forward(data[\"bert_input\"], data[\"segment_label\"])\n",
        "\n",
        "        next_loss = next_criterion(next_sent_output, data[\"is_next\"])\n",
        "        mask_loss = masked_criterion(mask_lm_output.transpose(1, 2), data[\"bert_label\"])\n",
        "        \n",
        "        #print(\"Next loss = \" + str(next_loss))\n",
        "        #print(\"Mask loss = \" + str(mask_loss))\n",
        "\n",
        "        loss = next_loss + mask_loss\n",
        "\n",
        "        correct = next_sent_output.argmax(dim=-1).eq(data[\"is_next\"]).sum().item()\n",
        "        \n",
        "        avg_next_loss += next_loss.item()\n",
        "        avg_mask_loss += mask_loss.item()\n",
        "        avg_loss += loss.item()\n",
        "        total_correct += correct\n",
        "        total_element += data[\"is_next\"].nelement()\n",
        "\n",
        "        #print(\"Epoch \" + str(l) + \" avg_loss = \" + str(avg_loss / (i + 1)) + \" total_acc = \" + str(total_correct * 100.0 / total_element))\n",
        "\n",
        "        optim_schedule.zero_grad()\n",
        "        loss.backward()\n",
        "        optim_schedule.step_and_update_lr()\n",
        "\n",
        "    train_loss.append(avg_loss)\n",
        "    train_acc.append(total_correct * 100.0 / total_element)\n",
        "    \n",
        "    print('Avg loss train ' + str(avg_loss / (i + 1)) + \" epoch \" + str(l))\n",
        "    print('Acc train ' + str(total_correct * 100.0 / total_element) + \" epoch \" +  str(l))\n",
        "   \n",
        "    tb.save_value('Avg loss train', 'avg_loss_train', l, avg_loss / (i + 1))\n",
        "    tb.save_value('Avg next loss train', 'avg_next_loss_train', l, avg_next_loss / (i + 1))\n",
        "    tb.save_value('Avg mask loss train', 'avg_mask_loss_train', l, avg_mask_loss / (i + 1))\n",
        "    tb.save_value('Acc train', 'acc_train', l, total_correct * 100.0 / total_element)\n",
        "\n",
        "    avg_loss = 0.0\n",
        "    avg_next_loss = 0.0\n",
        "    avg_mask_loss = 0.0\n",
        "    total_correct = 0\n",
        "    total_element = 0\n",
        "\n",
        "    print(\"Test\")\n",
        "\n",
        "    for i, data in enumerate(test_data_loader):\n",
        "        data = {key: value.to('cuda') for key, value in data.items()}\n",
        "\n",
        "        next_sent_output, mask_lm_output = model.forward(data[\"bert_input\"], data[\"segment_label\"])\n",
        "                                                                         \n",
        "        next_loss = next_criterion(next_sent_output, data[\"is_next\"])\n",
        "        mask_loss = masked_criterion(mask_lm_output.transpose(1, 2), data[\"bert_label\"])\n",
        "\n",
        "        loss = next_loss + mask_loss\n",
        "        loss = next_loss + mask_loss\n",
        "\n",
        "        correct = next_sent_output.argmax(dim=-1).eq(data[\"is_next\"]).sum().item()\n",
        "        \n",
        "        avg_next_loss += next_loss.item()\n",
        "        avg_mask_loss += mask_loss.item()\n",
        "        avg_loss += loss.item()\n",
        "        total_correct += correct\n",
        "        total_element += data[\"is_next\"].nelement()\n",
        "\n",
        "        #print(\"Epoch \" + str(l) + \" avg_loss = \" + str(avg_loss / (i + 1)) + \" total_acc = \" + str(total_correct * 100.0 / total_element))\n",
        "\n",
        "    test_loss.append(avg_loss)\n",
        "    test_acc.append(total_correct * 100.0 / total_element)\n",
        "    \n",
        "    print('Avg loss test ' + str(avg_loss / (i + 1)) + \" epoch \" + str(l))\n",
        "    print('Acc test ' + str(total_correct * 100.0 / total_element) + \" epoch \" +  str(l))\n",
        "    \n",
        "    tb.save_value('Avg loss test', 'avg_loss_test', l, avg_loss / (i + 1))\n",
        "    tb.save_value('Avg next loss test', 'avg_next_loss_test', l, avg_next_loss / (i + 1))\n",
        "    tb.save_value('Avg mask loss test', 'avg_mask_loss_test', l, avg_mask_loss / (i + 1))\n",
        "    tb.save_value('Acc test', 'acc_test', l, total_correct * 100.0 / total_element)\n",
        "    \n",
        "    #torch.save(model.cuda(), \"model_epoch_\" + str(l))\n",
        "    \n",
        "    print()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "6353\n",
            "Train\n",
            "Avg loss train 9.339153177185059 epoch 0\n",
            "Acc train 52.09788024470061 epoch 0\n",
            "Test\n",
            "Avg loss test 9.291124470703437 epoch 0\n",
            "Acc test 54.484 epoch 0\n",
            "\n",
            "Train\n",
            "Avg loss train 9.260304146728515 epoch 1\n",
            "Acc train 60.10615026537566 epoch 1\n",
            "Test\n",
            "Avg loss test 9.240381749389726 epoch 1\n",
            "Acc test 63.052 epoch 1\n",
            "\n",
            "Train\n",
            "Avg loss train 9.202615212402344 epoch 2\n",
            "Acc train 66.85941714854287 epoch 2\n",
            "Test\n",
            "Avg loss test 9.179325451326493 epoch 2\n",
            "Acc test 69.22 epoch 2\n",
            "\n",
            "Train\n",
            "Avg loss train 9.1598949899292 epoch 3\n",
            "Acc train 70.57167642919107 epoch 3\n",
            "Test\n",
            "Avg loss test 9.148575881558001 epoch 3\n",
            "Acc test 71.348 epoch 3\n",
            "\n",
            "Train\n",
            "Avg loss train 9.10768 epoch 4\n",
            "Acc train 74.3831859579649 epoch 4\n",
            "Test\n",
            "Avg loss test 9.082255093947701 epoch 4\n",
            "Acc test 75.705 epoch 4\n",
            "\n",
            "Train\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-31-967c9a7c80df>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    595\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    596\u001b[0m         \u001b[0moptim_schedule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 597\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    598\u001b[0m         \u001b[0moptim_schedule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep_and_update_lr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    599\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph)\u001b[0m\n\u001b[1;32m    105\u001b[0m                 \u001b[0mproducts\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mDefaults\u001b[0m \u001b[0mto\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    106\u001b[0m         \"\"\"\n\u001b[0;32m--> 107\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    108\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    109\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables)\u001b[0m\n\u001b[1;32m     91\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m     92\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 93\u001b[0;31m         allow_unreachable=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m     94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     95\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2ZcQWHAz88H6",
        "colab_type": "code",
        "outputId": "948f87ba-2512-4ca6-ff1c-0a4f57808694",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 322
        }
      },
      "source": [
        "!zip -r /content/Graph.zip /content/Graph\n",
        "\n",
        "from google.colab import files\n",
        "files.download(\"/content/Graph.zip\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "updating: content/Graph/ (stored 0%)\n",
            "updating: content/Graph/acc_train/ (stored 0%)\n",
            "updating: content/Graph/avg_loss_train/ (stored 0%)\n",
            "updating: content/Graph/avg_mask_loss_test/ (stored 0%)\n",
            "updating: content/Graph/avg_next_loss_test/ (stored 0%)\n",
            "updating: content/Graph/acc_test/ (stored 0%)\n",
            "updating: content/Graph/avg_next_loss_train/ (stored 0%)\n",
            "updating: content/Graph/avg_mask_loss_train/ (stored 0%)\n",
            "updating: content/Graph/avg_loss_test/ (stored 0%)\n",
            "  adding: content/Graph/acc_train/events.out.tfevents.1559358781.bad86040ca6c (deflated 48%)\n",
            "  adding: content/Graph/avg_loss_train/events.out.tfevents.1559358781.bad86040ca6c (deflated 52%)\n",
            "  adding: content/Graph/avg_mask_loss_test/events.out.tfevents.1559358808.bad86040ca6c (deflated 57%)\n",
            "  adding: content/Graph/avg_next_loss_test/events.out.tfevents.1559358808.bad86040ca6c (deflated 54%)\n",
            "  adding: content/Graph/acc_test/events.out.tfevents.1559358808.bad86040ca6c (deflated 51%)\n",
            "  adding: content/Graph/avg_next_loss_train/events.out.tfevents.1559358781.bad86040ca6c (deflated 55%)\n",
            "  adding: content/Graph/avg_mask_loss_train/events.out.tfevents.1559358781.bad86040ca6c (deflated 56%)\n",
            "  adding: content/Graph/avg_loss_test/events.out.tfevents.1559358808.bad86040ca6c (deflated 52%)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RpiOShbp0rb1",
        "colab_type": "code",
        "outputId": "2843054c-0f60-4b3a-edf2-0b74e6a3ff08",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 306
        }
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Sat Sep 14 11:13:49 2019       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 430.40       Driver Version: 418.67       CUDA Version: 10.1     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   73C    P8    35W / 149W |      0MiB / 11441MiB |      0%      Default |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                       GPU Memory |\n",
            "|  GPU       PID   Type   Process name                             Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TIC67efcffNX",
        "colab_type": "code",
        "outputId": "cdb672bf-7720-4c01-a329-a780b31e7b1d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 632
        }
      },
      "source": [
        "!pip3 install transformers"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting transformers\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/fd/f9/51824e40f0a23a49eab4fcaa45c1c797cbf9761adedd0b558dab7c958b34/transformers-2.1.1-py3-none-any.whl (311kB)\n",
            "\u001b[K     |████████████████████████████████| 317kB 6.5MB/s \n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from transformers) (4.28.1)\n",
            "Requirement already satisfied: boto3 in /usr/local/lib/python3.6/dist-packages (from transformers) (1.9.250)\n",
            "Collecting sacremoses (from transformers)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/1f/8e/ed5364a06a9ba720fddd9820155cc57300d28f5f43a6fd7b7e817177e642/sacremoses-0.0.35.tar.gz (859kB)\n",
            "\u001b[K     |████████████████████████████████| 860kB 44.0MB/s \n",
            "\u001b[?25hCollecting regex (from transformers)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/ff/60/d9782c56ceefa76033a00e1f84cd8c586c75e6e7fea2cd45ee8b46a386c5/regex-2019.08.19-cp36-cp36m-manylinux1_x86_64.whl (643kB)\n",
            "\u001b[K     |████████████████████████████████| 645kB 37.6MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from transformers) (1.16.5)\n",
            "Collecting sentencepiece (from transformers)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/14/3d/efb655a670b98f62ec32d66954e1109f403db4d937c50d779a75b9763a29/sentencepiece-0.1.83-cp36-cp36m-manylinux1_x86_64.whl (1.0MB)\n",
            "\u001b[K     |████████████████████████████████| 1.0MB 40.6MB/s \n",
            "\u001b[?25hRequirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from transformers) (2.21.0)\n",
            "Requirement already satisfied: jmespath<1.0.0,>=0.7.1 in /usr/local/lib/python3.6/dist-packages (from boto3->transformers) (0.9.4)\n",
            "Requirement already satisfied: botocore<1.13.0,>=1.12.250 in /usr/local/lib/python3.6/dist-packages (from boto3->transformers) (1.12.250)\n",
            "Requirement already satisfied: s3transfer<0.3.0,>=0.2.0 in /usr/local/lib/python3.6/dist-packages (from boto3->transformers) (0.2.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (1.12.0)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (7.0)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (0.14.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2019.9.11)\n",
            "Requirement already satisfied: idna<2.9,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2.8)\n",
            "Requirement already satisfied: urllib3<1.25,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (1.24.3)\n",
            "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (3.0.4)\n",
            "Requirement already satisfied: docutils<0.16,>=0.10 in /usr/local/lib/python3.6/dist-packages (from botocore<1.13.0,>=1.12.250->boto3->transformers) (0.15.2)\n",
            "Requirement already satisfied: python-dateutil<3.0.0,>=2.1; python_version >= \"2.7\" in /usr/local/lib/python3.6/dist-packages (from botocore<1.13.0,>=1.12.250->boto3->transformers) (2.5.3)\n",
            "Building wheels for collected packages: sacremoses\n",
            "  Building wheel for sacremoses (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for sacremoses: filename=sacremoses-0.0.35-cp36-none-any.whl size=883999 sha256=d23566e6c01906de8500c525009b0e8961bfc61885057a18d0a09b07a43b393b\n",
            "  Stored in directory: /root/.cache/pip/wheels/63/2a/db/63e2909042c634ef551d0d9ac825b2b0b32dede4a6d87ddc94\n",
            "Successfully built sacremoses\n",
            "Installing collected packages: sacremoses, regex, sentencepiece, transformers\n",
            "Successfully installed regex-2019.8.19 sacremoses-0.0.35 sentencepiece-0.1.83 transformers-2.1.1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rQoKOUEU0tCB",
        "colab_type": "code",
        "outputId": "59781c52-7567-43f2-c72c-42dc1e308fbc",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 136
        }
      },
      "source": [
        "import torch\n",
        "from transformers import *\n",
        "\n",
        "# Transformers has a unified API\n",
        "# for 8 transformer architectures and 30 pretrained weights.\n",
        "#          Model          | Tokenizer          | Pretrained weights shortcut\n",
        "MODELS = [\n",
        "          (DistilBertForMaskedLM, DistilBertTokenizer, 'distilbert-base-uncased')]\n",
        "\n",
        "# To use TensorFlow 2.0 versions of the models, simply prefix the class names with 'TF', e.g. `TFRobertaModel` is the TF 2.0 counterpart of the PyTorch model `RobertaModel`\n",
        "\n",
        "# Let's encode some text in a sequence of hidden-states using each model:\n",
        "for model_class, tokenizer_class, pretrained_weights in MODELS:\n",
        "    # Load pretrained model/tokenizer\n",
        "    tokenizer = tokenizer_class.from_pretrained(pretrained_weights)\n",
        "    model = model_class.from_pretrained(pretrained_weights)\n",
        "\n",
        "    # Encode text\n",
        "    input_ids = torch.tensor([tokenizer.encode(\"Here is some text to encode\", add_special_tokens=True)])  # Add special tokens takes care of adding [CLS], [SEP], <s>... tokens in the right way for each model.\n",
        "    with torch.no_grad():\n",
        "        last_hidden_states = model(input_ids)  # Models outputs are now tuples\n",
        "        print(last_hidden_states[0])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[[ -6.0730,  -6.0184,  -6.0532,  ...,  -5.3471,  -5.1327,  -3.2758],\n",
            "         [-11.2493, -11.0514, -11.2565,  ..., -10.8661,  -9.0513,  -7.8534],\n",
            "         [-12.1650, -11.9065, -12.0352,  ..., -10.5904,  -8.4455,  -7.9540],\n",
            "         ...,\n",
            "         [-12.6242, -12.2914, -12.5036,  ..., -11.2858, -10.2073, -11.0556],\n",
            "         [ -8.3184,  -8.1826,  -8.2792,  ...,  -7.9680,  -7.4850,  -8.3284],\n",
            "         [-11.3067, -11.3019, -11.3750,  ...,  -9.1290,  -9.4283,  -7.6095]]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lG7TBH0XyI3Y",
        "colab_type": "code",
        "outputId": "c8754711-fc4f-4eb6-81bd-a77b046f81d7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "len(last_hidden_states[0][0][0])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "30522"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FVGBZlQAf82F",
        "colab_type": "code",
        "outputId": "4a69e337-651f-49a4-a15e-10772a89effe",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "input_ids"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[  101,  2182,  2003,  2070,  3793,  2000,  4372, 16044,   102]])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 30
        }
      ]
    }
  ]
}